# -*- coding: UTF-8 -*-
from keras.models import Sequential,Model,load_model
from keras.layers import Convolution1D,Input
from keras.layers import MaxPool1D
from keras.layers import GRU,LSTM
from keras.layers import Dense
from keras.layers import Flatten
from keras.utils import plot_model
from keras.layers import TimeDistributed
from keras import optimizers
from DataTools import *
import numpy as np
import os


def ims_model_build(input_shape=(2000,1)):
    # input_shape = input.shape()
    cnn = Sequential()
    cnn.add(Convolution1D(10, 20, padding='same', activation='relu', input_shape=input_shape))
    cnn.add(MaxPool1D(pool_size=5))
    cnn.add(Convolution1D(20,10,padding='same',activation='relu'))
    cnn.add(MaxPool1D(pool_size=5))
    cnn.add(Flatten())
    cnn.add(Dense(20,activation='relu'))
    ims_model = Sequential()
    ims_model.add(TimeDistributed(cnn))
    # ims_model.add(LSTM(units=3, activation='sigmoid'))
    # ims_model.add(LSTM(units=32,return_sequences=True))
    # ims_model.add(LSTM(units=32,return_sequences=True))
    ims_model.add(GRU(32,return_sequences=False))
    ims_model.add(Dense(1, activation='sigmoid'))
    opt = optimizers.Adam(learning_rate=0.001)
    ims_model.compile(loss='mse', optimizer=opt, metrics=['accuracy'])
    # print(ims_model.summary())
    return ims_model

def visualize_loss(history, title="Training and Validation Loss"):
    loss = history.history["loss"]
    # val_loss = history.history["val_loss"]
    epochs = range(len(loss))
    plt.figure()
    plt.plot(epochs, loss, "b", label="Training loss")
    # plt.plot(epochs, val_loss, "r", label="Validation loss")
    plt.title(title)
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.savefig(title+'6.png')
    plt.show()

if __name__ == '__main__':
    # path = '../IMS/3rd_test/4th_test/txt'
    path = '../IMS/3rd_test/4th_test/txt'
    # data_sigma = imsdatasigma(path)
    # file = '2004.02.12.10.32.39'
    # filedir = os.path.join(path, file)
    datalist = getIMSdatalist(path)
    trainlist, test_list = ims_dataset_split(datalist)
    testdataset = [datalist[i] for i in test_list]
    traindataset = [datalist[k] for k in trainlist]
    train_generator = IMSDataGenerator(traindataset)
    test_generator = IMSDataGenerator(testdataset)
    # x = np.random.normal(size=(10, 3, 1, 2000))
    # x.shape
    # cnn = Sequential()
    # cnn.add(Convolution1D(10, 7, padding='same', activation='relu', input_dim=2000))
    # cnn.add(MaxPool1D(pool_size=2, strides=2))
    # cnn.add(Flatten())
    # cnn.add(Dense(1, activation='relu'))
    # # ims_model = Sequential()
    # # ims_model.add(TimeDistributed(cnn))
    # # ims_model.add(GRU(units=3, activation='relu'))
    # # ims_model.add(Dense(1, activation='sigmoid'))
    # cnn.compile(loss='MSE', optimizer='adam', metrics=['loss'])
    # # cnn.fit(test_generator)
    # print(cnn)
my_ims_model = load_model('my_ims_model-5.h5')
#my_ims_model = ims_model_build()
history = my_ims_model.fit(train_generator,verbose=1,epochs=10)
# plot_model(my_ims_model,to_file='my_ims_model.png')
my_ims_model.save('my_ims_model-5.h5')
test_age_predict = my_ims_model.predict(test_generator)
time = range(len(datalist))
age = [data[1] for data in datalist]

plt.figure()
plt.plot(time,age,'b',label='bearing_degration')
plt.scatter(test_list[:len(test_age_predict)],test_age_predict,c='r',marker='o')
plt.savefig('bearing_degration6.png')
plt.show()
visualize_loss(history)


